import openai, os
import pandas as pd
from openai.embeddings_utils import get_embedding, cosine_similarity
import openai, os, backoff

embedding_model = "text-embedding-ada-002"
openai.api_key = os.environ.get("OPENAI_API_KEY")


def recommend_product(df, product_name, n=3, pprint=True):
    product_embedding = df[df['product_name'] == product_name].iloc[0].embedding
    print(product_embedding)
    df["similarity"] = df.embedding.apply(lambda x: cosine_similarity(x, product_embedding))

    results = (
        df.sort_values("similarity", ascending=False)
        .head(n)
        .product_name
    )
    if pprint:
        for r in results:
            print(r)
    return results

df = pd.read_parquet("data/taobao_product_title.parquet")
results = recommend_product(df, "【新款】苹果Apple MacBook Pro 13", n=3)